package com.hzxgs.tracelogstarter.trace;

import com.hzxgs.tracelogstarter.properties.TraceProperties;
import com.hzxgs.tracelogstarter.util.TraceUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Correlation;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessagePostProcessor;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;

@Component
@ConditionalOnClass(value = { MessagePostProcessor.class })
@Slf4j
public class RabbitMQBeforeSendReplyPostProcessor implements MessagePostProcessor {

    @Resource
    private TraceProperties traceProperties;

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {
        return message;
    }

    @Override
    public Message postProcessMessage(Message message, Correlation correlation) {
        return MessagePostProcessor.super.postProcessMessage(message, correlation);
    }

    @Override
    public Message postProcessMessage(Message message, Correlation correlation, String exchange, String routingKey) {
        if (traceProperties.getEnable()) {
            try {
                MessageProperties messageProperties = message.getMessageProperties();
                String traceId = TraceUtils.getTraceId();
                // 不为空，说明当前线程有traceId
                if (!StringUtils.isBlank(traceId)) {
                    // 设置消息头
                    messageProperties.setHeader(TraceUtils.TRACE_ID_HEADER, traceId);
                    messageProperties.setHeader(TraceUtils.SPAN_ID_HEADER, TraceUtils.getNextSpanId());
                    return message;
                }
            } catch (Exception e) {
                log.error("rabbitMq设置traceId失败: {}", e.getMessage());
                e.printStackTrace();
            }
        }
        return MessagePostProcessor.super.postProcessMessage(message, correlation, exchange, routingKey);
    }
}
